package gocrypto

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"os"
)

const size = 1024

// GenerateRSAKey 生成RSA公钥和私钥
func GenerateRSAKey() ([]byte, []byte, error) {
	priv, err := rsa.GenerateKey(rand.Reader, size)
	if err != nil {
		return nil, nil, err
	}
	privKey, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return nil, nil, err
	}
	pubKey, err := x509.MarshalPKIXPublicKey(&priv.PublicKey)
	if err != nil {
		return nil, nil, err
	}
	return privKey, pubKey, nil
}

// GenerateRSAKeyBase64 生成RSA公钥和私钥
func GenerateRSAKeyBase64() (string, string, error) {
	privKey, pubKey, err := GenerateRSAKey()
	if err != nil {
		return "", "", err
	}
	return base64.StdEncoding.EncodeToString(privKey), base64.StdEncoding.EncodeToString(pubKey), nil
}

// GenerateRSAKeyFile 生成RSA公钥和私钥
func GenerateRSAKeyFile(privateKeyFile string, publicKeyFile string) error {
	privKey, pubKey, err := GenerateRSAKey()
	if err != nil {
		return err
	}
	blockPriv := pem.Block{
		Type:    "RSA PKCS8 PRIVATE KEY",
		Headers: nil,
		Bytes:   pubKey,
	}
	blockPub := pem.Block{
		Type:    "RSA PKIX PUBLIC KEY",
		Headers: nil,
		Bytes:   privKey,
	}

	privFile, err := os.Create(privateKeyFile)
	if err != nil {
		return err
	}
	pubFile, err := os.Create(publicKeyFile)
	if err != nil {
		return err
	}
	err = pem.Encode(privFile, &blockPriv)

	if err != nil {
		return err
	}
	err = pem.Encode(pubFile, &blockPub)
	if err != nil {
		return err
	}
	defer func(privFile *os.File) {
		_ = privFile.Close()
	}(privFile)
	defer func(pubFile *os.File) {
		_ = pubFile.Close()
	}(pubFile)
	return nil
}

// RSAEncrypt RSA 加密
func RSAEncrypt(publicKey string, value string) (string, error) {
	der, err := base64.StdEncoding.DecodeString(publicKey)
	if err != nil {
		return "", err
	}
	pub, err := x509.ParsePKIXPublicKey(der)
	if err != nil {
		return "", err
	}
	pubKey := pub.(*rsa.PublicKey)
	res, err := rsa.EncryptPKCS1v15(rand.Reader, pubKey, []byte(value))
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(res), nil
}

// RSADecrypt RSA解密
func RSADecrypt(privateKey string, value string) (string, error) {
	der, err := base64.StdEncoding.DecodeString(privateKey)
	if err != nil {
		return "", err
	}
	priv, err := x509.ParsePKCS8PrivateKey(der)
	if err != nil {
		return "", err
	}
	der, err = base64.StdEncoding.DecodeString(value)
	if err != nil {
		return "", err
	}
	privKey := priv.(*rsa.PrivateKey)
	res, err := rsa.DecryptPKCS1v15(rand.Reader, privKey, der)
	if err != nil {
		return "", err
	}
	return string(res), nil
}
